package assignment02;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;


import nlp.classify.BasicFeatureVector;
import nlp.classify.BasicLabeledFeatureVector;
import nlp.classify.FeatureExtractor;
import nlp.classify.FeatureVector;
import nlp.classify.LabeledFeatureVector;
import nlp.classify.LabeledInstance;
import nlp.classify.ProbabilisticClassifier;
import nlp.classify.ProbabilisticClassifierFactory;
import nlp.math.DoubleArrays;
import nlp.util.Counter;
import nlp.util.CounterMap;
import nlp.util.Indexer;

public class PerceptronClassifier<I, F, L> implements
		ProbabilisticClassifier<I, L> {
	private double[] weights;
	private Encoding<F, L> encoding;
	private IndexLinearizer indexLinearizer;
	private FeatureExtractor<I, F> featureExtractor;

	public PerceptronClassifier(double[] weights, Encoding<F, L> encoding,
			IndexLinearizer indexLinearizer,
			FeatureExtractor<I, F> featureExtractor) {
		this.weights = weights;
		this.encoding = encoding;
		this.indexLinearizer = indexLinearizer;
		this.featureExtractor = featureExtractor;
	}

	public static class Factory<I, F, L> implements
			ProbabilisticClassifierFactory<I, L> {
		FeatureExtractor<I, F> featureExtractor;

		public ProbabilisticClassifier<I, L> trainClassifier(
				List<LabeledInstance<I, L>> trainingData) {
			// build data encodings so the inner loops can be efficient
			Encoding<F, L> encoding = buildEncoding(trainingData);
			IndexLinearizer indexLinearizer = buildIndexLinearizer(encoding);
			double[] initialWeights = buildInitialWeights(indexLinearizer);
			EncodedDatum[] data = encodeData(trainingData, encoding);
			Optimizer<F, L> optimizer = new Optimizer<F, L>(encoding, data, indexLinearizer);
			double[] weights = optimizer.optimize(initialWeights);
			// build a classifer using these weights (and the data encodings)
			return new PerceptronClassifier<I, F, L>(weights, encoding,
					indexLinearizer, featureExtractor);
		}

		private double[] buildInitialWeights(IndexLinearizer indexLinearizer) {
			return DoubleArrays.constantArray(0.0,
					indexLinearizer.getNumLinearIndexes());
		}

		private IndexLinearizer buildIndexLinearizer(Encoding<F, L> encoding) {
			return new IndexLinearizer(encoding.getNumFeatures(),
					encoding.getNumLabels());
		}

		private Encoding<F, L> buildEncoding(List<LabeledInstance<I, L>> data) {
			// System.out.println("in Factory.buildEncoding(), data.size() " +
			// data.size());
			Indexer<F> featureIndexer = new Indexer<F>();
			Indexer<L> labelIndexer = new Indexer<L>();
			for (LabeledInstance<I, L> labeledInstance : data) {
				L label = labeledInstance.getLabel();
				Counter<F> features = featureExtractor
						.extractFeatures(labeledInstance.getInput());
				
				// System.out.println("features.size() is " + features.size());
				LabeledFeatureVector<F, L> labeledDatum = new BasicLabeledFeatureVector<F, L>(
						label, features);
				labelIndexer.add(labeledDatum.getLabel());
				for (F feature : labeledDatum.getFeatures().keySet()) {
					featureIndexer.add(feature);
					// System.out.println("feature is " + feature.toString());
				}
			}
			// System.out.println("featureIndexer size is " +
			// featureIndexer.size());
			return new Encoding<F, L>(featureIndexer, labelIndexer);
		}

		private EncodedDatum[] encodeData(List<LabeledInstance<I, L>> data, Encoding<F, L> encoding) {
			EncodedDatum[] encodedData = new EncodedDatum[data.size()];
			// System.out.println("in Factory.encodeData, data.size() is " +
			// data.size());
			for (int i = 0; i < data.size(); i++) {
				LabeledInstance<I, L> labeledInstance = data.get(i);
				L label = labeledInstance.getLabel();
				Counter<F> features = featureExtractor
						.extractFeatures(labeledInstance.getInput());
				LabeledFeatureVector<F, L> labeledFeatureVector = new BasicLabeledFeatureVector<F, L>(
						label, features);
				encodedData[i] = EncodedDatum.encodeLabeledDatum(
						labeledFeatureVector, encoding);
			}
			return encodedData;
		}
		
		public Factory(FeatureExtractor<I,F> featureExtractor) {
		     
		      this.featureExtractor = featureExtractor;
		}
	}
	
	
	
	public static class Optimizer<F, L> {
		private Encoding<F, L> encoding;
		private EncodedDatum[] data;
		private IndexLinearizer indexLinearizer;
		public Optimizer(Encoding<F, L> encoding, EncodedDatum[] data, IndexLinearizer indexLinearizer) {
			this.encoding = encoding;
			this.data = data;
			this.indexLinearizer = indexLinearizer;
		}
		
		double[] optimize(double[] weights) {
			double[] temp = Arrays.copyOf(weights, weights.length);
			double[] sums = Arrays.copyOf(weights, weights.length);
			System.out.println("weight.length is " + weights.length + ", and temp.length is " + temp.length);
			// TODO
			for (int i = 0; i < data.length; i++) {
				EncodedDatum datum = data[i];
				int numActiveFeatures = datum.getNumActiveFeatures();
				if (numActiveFeatures == 0) {
					System.out.println("what happened? datam number:	" + i);
				}
				Counter<Integer> counter = new Counter<Integer>();
				for (int labelIndex = 0; labelIndex < encoding.getNumLabels(); labelIndex++) {
					for (int feature = 0; feature < numActiveFeatures; feature++) {
						int featureIndex = datum.getFeatureIndex(feature);
						int linearIndex = indexLinearizer.getLinearIndex(featureIndex, labelIndex);
						
						counter.incrementCount(labelIndex, temp[linearIndex]);
					}
				}
//				temp = updateWeights(counter, datum.getLabelIndex(), temp);
				int expectedLabelIndex = datum.getLabelIndex();
				if (counter.size() == 0) {
					continue;
				}
				int predictedLabelIndex = counter.argMax();
				double expectedCount = counter.getCount(expectedLabelIndex);
				double predictedCount = counter.getCount(predictedLabelIndex);
//				System.out.println("datum.getLabelIndex() Is " + datum.getLabelIndex() + ", predictedLabelIndex is " + predictedLabelIndex);
				if (expectedCount < predictedCount) {
					
					for (int feature = 0; feature < numActiveFeatures; feature++) {
						int featureIndex = datum.getFeatureIndex(feature);
						int expectedLinearIndex = indexLinearizer.getLinearIndex(featureIndex,expectedLabelIndex);
						int predictedLinearIndex = indexLinearizer.getLinearIndex(featureIndex, predictedLabelIndex);
						temp[expectedLinearIndex] += 1.0;
						temp[predictedLinearIndex] -= 1.0;						
					}
				}
				// updateweights
				sums = updateSums(sums, temp);
			}
			temp = getAverage(sums);
			return temp;
		}
		
		private double[] updateSums(double[] sums, double[] temp) {
			for (int i = 0; i < sums.length; i++) {
				sums[i] += temp[i]; 
			}
			return sums;
		}
		private double[] getAverage(double[] sums) {
			for (int j = 0; j < sums.length; j++) {
				sums[j] /= data.length;
			}
			return sums;
		}
		
		
	}
	
	public static class IndexLinearizer {
	    int numFeatures;
	    int numLabels;

	    public int getNumLinearIndexes() {
	      return numFeatures * numLabels;
	    }

	    public int getLinearIndex(int featureIndex, int labelIndex) {
	      return labelIndex + featureIndex * numLabels;
	    }

	    public int getFeatureIndex(int linearIndex) {
	      return linearIndex / numLabels;
	    }

	    public int getLabelIndex(int linearIndex) {
	      return linearIndex % numLabels;
	    }

	    public IndexLinearizer(int numFeatures, int numLabels) {
	      this.numFeatures = numFeatures;
	      this.numLabels = numLabels;
	    }
	  }
	
	public static class EncodedDatum {

	    public static <F,L> EncodedDatum encodeDatum(FeatureVector<F> featureVector, Encoding<F, L> encoding) {
	      Counter<F> features = featureVector.getFeatures();
	      Counter<F> knownFeatures = new Counter<F>();
	      for (F feature : features.keySet()) {
	        if (encoding.getFeatureIndex(feature) < 0)
	          continue;
	        knownFeatures.incrementCount(feature, features.getCount(feature));
	      }
	      int numActiveFeatures = knownFeatures.keySet().size();
	      int[] featureIndexes = new int[numActiveFeatures];
	      double[] featureCounts = new double[knownFeatures.keySet().size()];
	      int i = 0;
	      for (F feature : knownFeatures.keySet()) {
	        int index = encoding.getFeatureIndex(feature);
	        double count = knownFeatures.getCount(feature);
	        featureIndexes[i] = index;
	        featureCounts[i] = count;
	        i++;
	      }
	      EncodedDatum encodedDatum = new EncodedDatum(-1, featureIndexes, featureCounts);
	      return encodedDatum;
	    }

	    public static <F,L> EncodedDatum encodeLabeledDatum(LabeledFeatureVector<F, L> labeledDatum, Encoding<F, L> encoding) {
	      EncodedDatum encodedDatum = encodeDatum(labeledDatum, encoding);
	      encodedDatum.labelIndex = encoding.getLabelIndex(labeledDatum.getLabel());
	      return encodedDatum;
	    }

	    int labelIndex;
	    int[] featureIndexes;
	    double[] featureCounts;

	    public int getLabelIndex() {
	      return labelIndex;
	    }

	    public int getNumActiveFeatures() {
	      return featureCounts.length;
	    }

	    public int getFeatureIndex(int num) {
	      return featureIndexes[num];
	    }

	    public double getFeatureCount(int num) {
	      return featureCounts[num];
	    }

	    public EncodedDatum(int labelIndex, int[] featureIndexes, double[] featureCounts) {
	      this.labelIndex = labelIndex;
	      this.featureIndexes = featureIndexes;
	      this.featureCounts = featureCounts;
	    }
	  }
	
	public static class Encoding <F,L> {
	    Indexer<F> featureIndexer;
	    Indexer<L> labelIndexer;

	    public int getNumFeatures() {
	      return featureIndexer.size();
	    }

	    public int getFeatureIndex(F feature) {
	      return featureIndexer.indexOf(feature);
	    }

	    public F getFeature(int featureIndex) {
	      return featureIndexer.get(featureIndex);
	    }

	    public int getNumLabels() {
	      return labelIndexer.size();
	    }

	    public int getLabelIndex(L label) {
	      return labelIndexer.indexOf(label);
	    }

	    public L getLabel(int labelIndex) {
	      return labelIndexer.get(labelIndex);
	    }

	    public Encoding(Indexer<F> featureIndexer, Indexer<L> labelIndexer) {
	      this.featureIndexer = featureIndexer;
	      this.labelIndexer = labelIndexer;
	    }
	  }

	public L getLabel(I instance) {
		return getProbabilities(instance).argMax();
	}

	private static <F, L> double[] getLogProbabilities(EncodedDatum datum,
			double[] weights, Encoding<F, L> encoding,
			IndexLinearizer indexLinearizer) {
		// TODO: apply the classifier to this feature vector
		// dummy code
//		 double[] logProbabilities =
//		 DoubleArrays.constantArray(Double.NEGATIVE_INFINITY, encoding.getNumLabels());
//		 logProbabilities[0] = 0.0;
//		 return logProbabilities;
		// end dummy code

		// TODO
		double[] logProbabilities = DoubleArrays.constantArray(Double.NEGATIVE_INFINITY, encoding.getNumLabels());
		CounterMap<EncodedDatum, Integer> pHistoryMap = new CounterMap<EncodedDatum, Integer>();
		int numOfLabels = indexLinearizer.numLabels;
		// for (int i = 0; i < 10; i++) {
		// System.out.println("data[" + i + "] is " + data[i]);
		// }
		Counter<Integer> datumCounter = pHistoryMap.getCounter(datum);
		int featureIndexes[] = datum.featureIndexes;
		for (int j = 0; j < numOfLabels; j++) {
			int labelIndex = j;
			double sum = 0.0;
			for (int k = 0; k < featureIndexes.length; k++) {
				int featureIndex = featureIndexes[k];
				int linearIndex = indexLinearizer.getLinearIndex(featureIndex,
						labelIndex);
				double weight = weights[linearIndex];
				sum += weight * 1.0;
			}
			datumCounter.setCount(labelIndex, sum);
			// datumCounter.incrementCount(labelIndex, sum);
		}
//		System.out.println("initializing countermap done.");

		// normalize logProbabilities
		for (Integer labelIndex : datumCounter.keySet()) {
			double count = datumCounter.getCount(labelIndex);
			datumCounter.setCount(labelIndex, Math.exp(count));
		}
		double denominator = datumCounter.totalCount();
		int i = 0;
		for (Integer labelIndex : datumCounter.keySet()) {
			double count = datumCounter.getCount(labelIndex);
			datumCounter.setCount(labelIndex, count / denominator);
			logProbabilities[i] = Math.log(count / denominator);
//			System.out.println("logProbabilities is " + logProbabilities[i]);
			i++;
		}
//		System.out.println("calculating probability of each pair of datum and label done");

		return logProbabilities;
	}

  public Counter<L> getProbabilities(I input) {
    FeatureVector<F> featureVector = new BasicFeatureVector<F>(featureExtractor.extractFeatures(input));
    return getProbabilities(featureVector);
  }

  private Counter<L> getProbabilities(FeatureVector<F> featureVector) {
    EncodedDatum encodedDatum = EncodedDatum.encodeDatum(featureVector, encoding);
    double[] logProbabilities = getLogProbabilities(encodedDatum, weights, encoding, indexLinearizer);
    return logProbabiltyArrayToProbabiltyCounter(logProbabilities);
  }

  private Counter<L> logProbabiltyArrayToProbabiltyCounter(double[] logProbabilities) {
    Counter<L> probabiltyCounter = new Counter<L>();
    for (int labelIndex = 0; labelIndex < logProbabilities.length; labelIndex++) {
      double logProbability = logProbabilities[labelIndex];
      double probability = Math.exp(logProbability);
      L label = encoding.getLabel(labelIndex);
      probabiltyCounter.setCount(label, probability);
    }
    return probabiltyCounter;
  }

  public static void main(String[] args) {
		// create datums
		LabeledInstance<String[], String> datum1 = new LabeledInstance<String[], String>("cat", new String[] { "fuzzy", "claws", "small" });
		LabeledInstance<String[], String> datum2 = new LabeledInstance<String[], String>("bear", new String[] { "fuzzy", "claws", "big" });
		LabeledInstance<String[], String> datum3 = new LabeledInstance<String[], String>("cat", new String[] { "claws", "medium" });
		LabeledInstance<String[], String> datum4 = new LabeledInstance<String[], String>("cat", new String[] { "claws", "small" });

		// create training set
		List<LabeledInstance<String[], String>> trainingData = new ArrayList<LabeledInstance<String[], String>>();
		trainingData.add(datum1);
		trainingData.add(datum2);
		trainingData.add(datum3);

		// create test set
		List<LabeledInstance<String[], String>> testData = new ArrayList<LabeledInstance<String[], String>>();
		testData.add(datum4);

		// build classifier
		FeatureExtractor<String[], String> featureExtractor = new FeatureExtractor<String[], String>() {
			public Counter<String> extractFeatures(String[] featureArray) {
				return new Counter<String>(Arrays.asList(featureArray));
			}
		};
		PerceptronClassifier.Factory<String[], String, String> perceptronClassifierFactory = 
			new PerceptronClassifier.Factory<String[], String, String>(featureExtractor);
		ProbabilisticClassifier<String[], String> perceptronClassifier = perceptronClassifierFactory
				.trainClassifier(trainingData);
		
		LabeledInstance<String[], String> datum555 = new LabeledInstance<String[], String>("cat", new String[]{"fuzzy", "fuzzy", "big"});
		System.out.println("Probabilities on test instance: "
				+ perceptronClassifier.getProbabilities(datum4.getInput()));
	}
}
